import json
import os
from nodes import CFG_Node, CDG_Node, CG_Node, DDG_Node, Edge

def json2node(directory_path):
    nodes = []
    edges = []

    # 映射文件名到相应的节点类型
    node_type_map = {
        'cfg_graph.json': CFG_Node,
        'cdg_graph.json': CDG_Node,
        'ddg_graph.json': DDG_Node,
        'cg_graph.json': CG_Node
    }

    # 遍历指定目录下的文件
    for file_name in os.listdir(directory_path):
        if file_name in node_type_map:
            file_path = os.path.join(directory_path, file_name)
            node_class = node_type_map[file_name]

            # 读取JSON文件
            with open(file_path, 'r') as f:
                data = json.load(f)

            # 解析节点信息
            for node_info in data.get('nodes', []):
                try:
                    node_obj = node_class(
                        id=node_info['id'],
                        addr=node_info['addr'],
                        function_name=node_info.get('function_name'),
                        function_addr=node_info.get('function_addr'),
                        block_addr=node_info.get('block_addr'),
                        disasm=node_info.get('disasm')
                    )
                    nodes.append(node_obj)
                except Exception as e:
                    print(f"Error processing node {node_info.get('addr', 'unknown')} in {file_name}: {e}")

            # 解析边信息
            for edge_info in data.get('edges', []):
                try:
                    edge_obj = Edge(
                        src_id=edge_info['src_id'],
                        dst_id=edge_info['dst_id'],
                        type=edge_info.get('type')  
                    )
                    edges.append(edge_obj)
                except Exception as e:
                    print(f"Error processing edge from {edge_info.get('src_id', 'unknown')} to {edge_info.get('dst_id', 'unknown')} in {file_name}: {e}")

    return nodes, edges

# 示例用法
if __name__ == "__main__":
    directory_path = '/home/kali/Desktop/angr/graph/CWE121_Stack_Based_Buffer_Overflow/s01/bad/CWE121_Stack_Based_Buffer_Overflow__char_type_overrun_memcpy_01'
    nodes, edges = json2node(directory_path)

    # 打印读取到的节点和边的信息
    for node in nodes:
        print(node)

    for edge in edges:
        print(edge)
